#include <cstddef>
#include <cuda_runtime.h>
#include <cuda_fp16.h>
#include <numeric>
#include <vector>
#include<fstream>
#include<iostream>
#include<tuple>

#include "common/tensorrt.hpp"
#include "common/check.hpp"
#include "cuda_util.h"
#include "bev_model.hpp"


namespace dsgn{
    class BevImpl : public BevModel{
        public:
            virtual ~BevImpl() {
                if (bev_feat_) checkRuntime(cudaFree(bev_feat_));
            }
            
            bool init(const std::string& model) {
                engine_ = TensorRT::load(model);
                if (engine_ == nullptr) return false;
                std::vector<int> cls_size_dims = engine_->static_dims("cls_preds");
                cls_size_ = std::accumulate(cls_size_dims.begin(), cls_size_dims.end(), 1, std::multiplies<int>())*sizeof(float);
                std::vector<int> box_size_dims = engine_->static_dims("box_preds");
                box_size_ = std::accumulate(box_size_dims.begin(), box_size_dims.end(), 1, std::multiplies<int>())*sizeof(float);
                std::vector<int> dir_cls_size_dims = engine_->static_dims("dir_cls_preds");
                dir_cls_size_ = std::accumulate(dir_cls_size_dims.begin(), dir_cls_size_dims.end(), 1, std::multiplies<int>())*sizeof(float);


                checkRuntime(cudaMallocManaged((void **)&dir_cls_preds_, dir_cls_size_));
                // printf("dir_cls_preds_ size: %d\n", dir_cls_size_);
                checkRuntime(cudaMallocManaged((void **)&box_preds_, box_size_));
                // printf("box_preds_ size: %d\n", box_size_);
                checkRuntime(cudaMallocManaged((void **)&cls_preds_, cls_size_));
                // printf("cls_preds_ size: %d\n", cls_size_);

                // std::vector<int> batch_cls_size_dims = engine_->static_dims("batch_cls_preds");
                // unsigned int batch_cls_size_ = std::accumulate(batch_cls_size_dims.begin(), batch_cls_size_dims.end(), 1, std::multiplies<int>())*sizeof(float);
                // std::vector<int> batch_box_size_dims = engine_->static_dims("batch_box_preds");
                // unsigned int batch_box_size_ = std::accumulate(batch_box_size_dims.begin(), batch_box_size_dims.end(), 1, std::multiplies<int>())*sizeof(float);
                // checkRuntime(cudaMallocManaged((void **)&batch_cls_preds_, batch_cls_size_));
                // // printf("batch_cls_preds size: %d\n", batch_cls_size_);
                // checkRuntime(cudaMallocManaged((void **)&batch_box_preds_, batch_box_size_));

                return true;
            }

            virtual void print() override { engine_->print("Feature Extractor Backbone"); }

            virtual void forward(const float* images,
                        const float* downsampled_disp, const float* norm_coord_imgs,
                        const float* valids, //const float* valids2d, 
                        // const int* psv_disps_channels, const int* voxel_disps_channels, 
                        void* stream=nullptr) override {

                engine_->forward({images, downsampled_disp, norm_coord_imgs, valids, // valids2d, psv_disps_channels, voxel_disps_channels,  // Input
                                 box_preds_, cls_preds_, dir_cls_preds_,
                                //  batch_cls_preds_, batch_box_preds_
                                 },  // Output},
                                 stream);

            }

            virtual float* get_bev_feature() override { return bev_feat_; }
            virtual float* get_dir_cls_preds() override { return dir_cls_preds_; }
            virtual float* get_box_preds() override { return box_preds_; }
            virtual float* get_cls_preds() override { return cls_preds_; }


        private:
            std::shared_ptr<TensorRT::Engine> engine_;
            size_t bytes_image_;
            // std::vector<half*> bindings_;
            // std::vector<std::vector<int>> bindshape_;
            float *dir_cls_preds_ = nullptr;
            float *box_preds_ = nullptr;
            float *cls_preds_ = nullptr;
            float *batch_cls_preds_ = nullptr;
            float *batch_box_preds_ = nullptr;
            unsigned int cls_size_;
            unsigned int box_size_;
            unsigned int dir_cls_size_;
            float *bndbox_output_ = nullptr;
            unsigned int bndbox_size_ = 0;
            // std::vector<Box> res_;

            float* bev_feat_ = nullptr;
    };

    std::shared_ptr<BevModel> create_bevmodel(const std::string& model) {
        std::shared_ptr<BevImpl> instance(new BevImpl());
        if (!instance->init(model)){
            instance.reset();
        }
        return instance;
    };
} // namespace dsgn
